import cv2
import numpy as np
import os

import matplotlib.pyplot as plt
import rasterio
import tarfile

from matplotlib.colors import LinearSegmentedColormap
from PIL import Image
from rasterio.io import MemoryFile


def read_etopo_tiles():
    cmap = LinearSegmentedColormap.from_list('evaluation', {
        0: [2, 12, 43], 20: [19, 39, 20]
    })
    parent_directory = 'E:\\map data'
    file = 'ETOPO2022.tar'
    directory_name = os.path.join(parent_directory, file)

    all_width = 7200; all_height = 3600
    all_images = np.zeros((all_height, all_width)) #row, column
    single_width = int(7200 / (360 / 15)); single_height = int(3600 / (180 / 15))

    if directory_name.endswith('.tar'):
        tar = tarfile.open(directory_name)

        tar_members = tar.getmembers()
        for tar_member in tar_members:
            tar_member_name = tar_member.name
            if tar_member.name.endswith('.tif'):
                extracted_tif_file = tar.extractfile(tar_member)
                content = extracted_tif_file.read()
                # tiff_bytes = io.BytesIO(content)
                print(tar_member_name)

                with MemoryFile(content) as memfile:
                    # 打开内存文件中的 TIFF 数据集
                    with memfile.open() as dataset:
                        # 示例：读取影像数据（numpy 数组）
                        image_data = dataset.read()  # shape: (波段数, 高度, 宽度)

                        # 示例：获取元信息（如投影、变换参数等）
                        width = dataset.width; height = dataset.height
                        bounds = dataset.bounds
                        meta = dataset.meta
                        crs = dataset.crs
                        transform = dataset.transform

                        lon_range = bounds.right - bounds.left
                        lat_range = bounds.bottom - bounds.top

                        transformer = rasterio.transform.AffineTransformer(transform)
                        geo_point = transformer.xy(0, 0)
                        print(f"影像形状: {image_data.shape}")
                        print(f"投影信息: {crs}")

                        image_nparray = image_data[0]
                        # image_nparray[image_nparray < 0] = 10000
                        resized_image_cv2 = cv2.resize(image_nparray, (single_width, single_height), interpolation=cv2.INTER_AREA)

                        left_ratio = (bounds.left - (-180)) / (180 - (-180))
                        top_ratio = 1 - (bounds.top - (-90)) / (90 - (-90))

                        row_start = round(top_ratio * all_height); column_start = round(left_ratio * all_width)

                        print(f"row_start: {row_start}; column_start: {column_start}")
                        all_images[row_start:(row_start + single_height), column_start:(column_start + single_width)] = resized_image_cv2

        tar.close()

    plt.imshow(all_images)
    plt.show()

    image = Image.fromarray(all_images)
    image.show()
                        # 如需返回数据，可根据需求返回 dataset、image_data 或 meta 等
                        # return dataset  # 注意：若在外部使用，需确保内存文件未关闭（建议在 with 块内处理）

    return 0

if __name__ == '__main__':
    read_etopo_tiles()